package com.snail.generator.config;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.snail.generator.common.WebContext;
import com.snail.generator.entity.DataSourceInfo;
import com.snail.generator.utils.SqlUtils;
import com.snail.generator.mapper.GeneratorMapper;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.HandlerInterceptor;

import java.io.BufferedReader;
import java.sql.Connection;
import java.sql.DriverManager;

@Component
@RequiredArgsConstructor
public class DataSourceConnectInterceptor implements HandlerInterceptor {

    private final ObjectMapper objectMapper;
    private final GeneratorMapper[] generatorMappers;

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        // 读取请求体中的内容
        BufferedReader reader = request.getReader();
        StringBuilder requestBody = new StringBuilder();
        String line;
        while ((line = reader.readLine()) != null) {
            requestBody.append(line);
        }
        String body = requestBody.toString();

        DataSourceInfo dataSourceInfo = objectMapper.readValue(body, DataSourceInfo.class);

        String database = request.getParameter("database");
        if (database == null) database = "";

        Class.forName(SqlUtils.getDriverClassName(dataSourceInfo));
        Connection connection = DriverManager.getConnection(
                SqlUtils.getUrl(dataSourceInfo) + database,
                dataSourceInfo.getUsername(),
                dataSourceInfo.getPassword()
        );
        connection.setAutoCommit(false);

        Class<?> mapperClass = SqlUtils.getMapperClass(dataSourceInfo);
        for (GeneratorMapper mapper : generatorMappers) {
            if (mapperClass.equals(mapper.getClass())) {
                WebContext.setGeneratorMapper(mapper);
                break;
            }
        }

        WebContext.setDataSourceConnection(connection);

        return true;
    }

    @Override
    public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, Exception ex) throws Exception {
        Connection connection = WebContext.getDataSourceConnection();
        if (connection == null) return;
        try (connection) {
            if (ex != null) connection.rollback();
            else connection.commit();
        } finally {
            WebContext.removeDataSourceConnection();
        }
    }
}
